#include "../Base.h"


NTSTATUS NTAPI NtGetNextThread(
	__in HANDLE ProcessHandle,
	__in HANDLE ThreadHandle,
	__in ACCESS_MASK DesiredAccess,
	__in ULONG HandleAttributes,
	__in ULONG Flags,
	__out PHANDLE NewThreadHandle
)
{

	typedef NTSTATUS(NTAPI* ZwGetNextThreadProc)(
		__in HANDLE ProcessHandle,
		__in HANDLE ThreadHandle,
		__in ACCESS_MASK DesiredAccess,
		__in ULONG HandleAttributes,
		__in ULONG Flags,
		__out PHANDLE NewThreadHandle
		);

	static ZwGetNextThreadProc ZwGetNextThreadFunc = NULL;
	if (!ZwGetNextThreadFunc)
	{
		UNICODE_STRING unName = { 0 };
		RtlInitUnicodeString(&unName, L"ZwGetNextThread");
		ZwGetNextThreadFunc = (ZwGetNextThreadProc)MmGetSystemRoutineAddress(&unName);
		if (!ZwGetNextThreadFunc)
		{
			UNICODE_STRING unName = { 0 };
			RtlInitUnicodeString(&unName, L"ZwGetNotificationResourceManager");
			PUCHAR ZwGetNotificationResourceManagerAddr = (PUCHAR)MmGetSystemRoutineAddress(&unName);
			ZwGetNotificationResourceManagerAddr -= 0x50;
			for (int i = 0; i < 0x30; i++)
			{
				if (ZwGetNotificationResourceManagerAddr[i] == 0x48
					&& ZwGetNotificationResourceManagerAddr[i + 1] == 0x8B
					&& ZwGetNotificationResourceManagerAddr[i + 2] == 0xC4)
				{
					ZwGetNextThreadFunc = (ZwGetNextThreadProc)(ZwGetNotificationResourceManagerAddr + i);
					break;
				}
			}
		}
	}

	if (ZwGetNextThreadFunc)
	{
		return ZwGetNextThreadFunc(ProcessHandle, ThreadHandle, DesiredAccess,
			HandleAttributes, Flags, NewThreadHandle);
	}

	return STATUS_UNSUCCESSFUL;
}


PETHREAD NtGetProcessMainThread(PEPROCESS Process)
{
	PETHREAD ethread = NULL;

	KAPC_STATE kApcState = { 0 };

	KeStackAttachProcess(Process, &kApcState);

	HANDLE hThread = NULL;

	NTSTATUS status = NtGetNextThread(NtCurrentProcess(), NULL, THREAD_ALL_ACCESS,
		OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, 0, &hThread);

	if (NT_SUCCESS(status))
	{

		status = ObReferenceObjectByHandle(hThread, THREAD_ALL_ACCESS,
			*PsThreadType, KernelMode, (PVOID*) & ethread, NULL);
		NtClose(hThread);

		if (!NT_SUCCESS(status))
		{
			ethread = NULL;
		}
	}


	KeUnstackDetachProcess(&kApcState);
	return ethread;
}

NTSTATUS
PsSuspendThread(
	IN PETHREAD Thread,
	OUT PULONG PreviousSuspendCount
) {

	typedef NTSTATUS
	(*pfnPsSuspendThread)(
		IN PETHREAD Thread,
		OUT PULONG PreviousSuspendCount
		);
	static pfnPsSuspendThread PsSuspendThreadProc = NULL;
	if (!PsSuspendThreadProc) {
		PsSuspendThreadProc = (pfnPsSuspendThread) searchCode("ntoskrnl.exe", "PAGE", "48***4C8BF2488BF98364***65********4889***66FF*****4C8D*****4C89***498BCFE8****",-0x11);

		

	}
	if (PsSuspendThreadProc) {
		PsSuspendThreadProc(Thread, PreviousSuspendCount);
	}

	return STATUS_NOT_IMPLEMENTED;


}

NTSTATUS
PsResumeThread(
	IN PETHREAD Thread,
	OUT PULONG PreviousSuspendCount OPTIONAL
) {

	
	typedef NTSTATUS
	(*pfnPsResumeThread)(
		IN PETHREAD Thread,
		OUT PULONG PreviousSuspendCount OPTIONAL
		);
	static pfnPsResumeThread PsResumeThreadProc = NULL;
	if (!PsResumeThreadProc) {
		PsResumeThreadProc = (pfnPsResumeThread)searchCode("ntoskrnl.exe", "PAGE", "48***488BDA488BF9E8****65********8BF083**75*4C8B*****B8****418B*****", -0xb);



	}
	if (PsResumeThreadProc) {
		PsResumeThreadProc(Thread, PreviousSuspendCount);
	}

	return STATUS_NOT_IMPLEMENTED;
}



ULONG GetThreadStartAddressOffset() {
	static ULONG offset = 0;
	if (offset) return offset;
	RTL_OSVERSIONINFOW info;
	RtlGetVersion(&info);

	UNICODE_STRING uName = { 0 };
	RtlInitUnicodeString(&uName, L"PsGetThreadId");
	PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&uName);
	ULONG tidOffset = 0;
	if (!func) return 0;

	for (int i = 0; i < 100; i++) {
		if (func[i] == 0xc3 && (func[i + 1] == 0xcc || func[i + 1] == 0x90) && (func[i + 2] == 0xcc || func[i + 2] == 0x90)) {

			tidOffset = *(PULONG)(func + i - 4);
			break;
		}
	}

	if (!tidOffset) return 0;

	

	offset = tidOffset - 0x30;


	return offset;
}


ULONG GetWin32ThreadStartAddressOffset() {
	static ULONG offset = 0;
	if (offset) return offset;
	RTL_OSVERSIONINFOW info;
	RtlGetVersion(&info);

	UNICODE_STRING uName = { 0 };
	RtlInitUnicodeString(&uName, L"PsGetThreadId");
	PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&uName);
	ULONG tidOffset = 0;
	if (!func) return 0;

	for (int i = 0; i < 100; i++) {
		if (func[i] == 0xc3 && (func[i + 1] == 0xcc || func[i + 1] == 0x90) && (func[i + 2] == 0xcc || func[i + 2] == 0x90)) {

			tidOffset = *(PULONG)(func + i  - 4);
			break;
		}
	}

	if (!tidOffset) return 0;


	if (info.dwBuildNumber == 7600 || info.dwBuildNumber == 7601) {

		offset = tidOffset + 0x58;
	}
	else {
		offset = tidOffset + 0x50;
	}

	return offset;


}


ULONG GetEThreadListOffset() {
	static ULONG offset = 0;
	if (offset) return offset;
	RTL_OSVERSIONINFOW info;
	RtlGetVersion(&info);

	UNICODE_STRING uName = { 0 };
	RtlInitUnicodeString(&uName, L"PsGetThreadId");
	PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&uName);
	ULONG tidOffset = 0;
	if (!func) return 0;

	for (int i = 0; i < 100; i++) {
		if (func[i] == 0xc3 && (func[i + 1] == 0xcc || func[i + 1] == 0x90) && (func[i + 2] == 0xcc || func[i + 2] == 0x90)) {

			tidOffset = *(PULONG)(func + i - 4);
			break;
		}
	}

	if (!tidOffset) return 0;


	if (info.dwBuildNumber == 7600 || info.dwBuildNumber == 7601 || info.dwBuildNumber >=16299) {

		offset = tidOffset + 0x68;
	}
	else {
		offset = tidOffset + 0x60;
	}

	return offset;


}